import torch.nn as nn

from utils import clone_module
from sublayer_connection import SublayerConnection


class DecoderLayer(nn.Module):

    def __init__(self, feature_dim, self_attn, src_attn, feed_forward, dropout):
        """
        size，代表词嵌入的维度大小, 同时也代表解码器层的尺寸，
        self_attn: 多头自注意力对象 Q=K=V， 
        src_attn: 多头注意力对象，这里Q!=K=V
        feed_forward: 前馈全连接层实例化对象
        """
        super(DecoderLayer, self).__init__()
        self.feature_dim = feature_dim
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayers = clone_module(SublayerConnection(feature_dim, dropout), 3)

    def forward(self, x, memory, source_mask, target_mask):
        """
        x: 上一层的输出
        mermory: 来自编码器的语义存储变量 
        source_mask: 源数据掩码张量
        target_mask: 目标数据掩码张量.
        """

        # 将x传入第一个子层结构，第一个子层结构的输入分别是x和self-attn函数，因为是自注意力机制，所以Q,K,V都是x，
        # 最后一个参数是目标数据掩码张量，这时要对目标数据进行遮掩，因为此时模型可能还没有生成任何目标数据，
        # 比如在解码器准备生成第一个字符或词汇时，我们其实已经传入了第一个字符以便计算损失，
        # 但是我们不希望在生成第一个字符时模型能利用这个信息，因此我们会将其遮掩，同样生成第二个字符或词汇时，
        # 模型只能使用第一个字符或词汇信息，第二个字符以及之后的信息都不允许被模型使用.
        x = self.sublayers[0](x, lambda x: self.self_attn(x, x, x, target_mask))

        # 接着进入第二个子层，这个子层中常规的注意力机制，q是输入x; k，v是编码层输出memory，
        # 同样也传入source_mask，但是进行源数据遮掩的原因并非是抑制信息泄漏，而是遮蔽掉对结果没有意义的字符而产生的注意力值，
        # 以此提升模型效果和训练速度. 这样就完成了第二个子层的处理.
        x = self.sublayers[1](x, lambda x: self.src_attn(x, memory, memory, source_mask))

        return self.sublayers[2](x, self.feed_forward)


if __name__ == '__main__':
    import torch
    import copy

    from multi_headed_attention import MultiHeadedAttention
    from positionwise_feed_forward import PositionwiseFeedForward
    from utils import subsequent_mask
    from utils import clone_module

    batch_size = 3
    seq_len = 4
    feature_dim = 6
    num_head = 2
    d_fff = 54
    dropout = 0.2

    self_attn = MultiHeadedAttention(
        num_head=num_head,
        feature_dim=feature_dim,
        dropout=dropout,
    )

    src_attn = clone_module(self_attn, 1)[0]

    feed_forward = PositionwiseFeedForward(
        feature_dim=feature_dim,
        d_ff=d_fff,
        dropout=dropout,
    )

    decoder_layer = DecoderLayer(
        feature_dim=feature_dim,
        self_attn=self_attn,
        src_attn=src_attn,
        feed_forward=feed_forward,
        dropout=dropout,
    )

    target_mask = subsequent_mask(size=seq_len)
    source_mask = copy.deepcopy(target_mask)  # 实际中source_mask和target_mask并不相同
    x = torch.randn(batch_size, seq_len, feature_dim)
    memory = torch.randn(batch_size, seq_len, feature_dim)

    result = decoder_layer(
        x=x,
        memory=memory,
        source_mask=source_mask,
        target_mask=target_mask,
    )

    print('x size: ', x.size())  # [batch_size, seq_len, feature_dim]
    print('result size: ', result.size())  # [batch_size, seq_len, feature_dim]
    print(result)
